
from typing import Dict, Iterator, List, TypedDict

from app.jobs.pdfparser.internal.logic.splitter.llm.paper_abstract_extractor import AbstractState
from langchain_openai import ChatOpenAI
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langgraph.config import get_stream_writer
from langgraph.graph import StateGraph, START, END

from .by_3steps_splitter_utils import *
from .paper_abstract_extractor_utils import merge_paper_basic_info

import sys
def cout(s: str):
    print(s, file=sys.stderr)

class SplitState(TypedDict):
    paper_md_list: List[str]
    
    index: int
    text: str

    paper_info: str
    abstract_summary: Dict

    split_tip: str
    split_res: str
    next_tip: str
    
    next_or_end: str

class By3StepsSplitter():
    """
    三步法解析和分割论文文本：
    第一步 —— 文本段结构分析与分割建议生成
    第二步 —— 文本分割与公式 / 表格纠正
    第三步 —— 下一段分割铺垫信息提取
    """
    
    first_tip = "这是论文中的第一段分割文本，还没有对上一段文本的分析"
    
    def __init__(self, llm: ChatOpenAI, max_chunk_size: int = 2048):
        self._llm = llm
        self._text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=980,
            chunk_overlap=80,
            length_function=len,
            separators=["\n\n",
                        ".", "。", "\u3002",
                        ",", "，", "\uff0c",
                        "\n",
                        " ", "\u3000", "\u200b",    # 空格/全角空格/零宽空格
                        ''],
            is_separator_regex=False,
        )
        self._max_chunk_size = max_chunk_size
        
        self._build_graph()
        
    def split_stream(self, abstract_state: AbstractState) -> Iterator[str]:
        abstract_summary = abstract_state["my_summary"]
        paper_info = merge_paper_basic_info(abstract_summary)
        paper_md_list = abstract_state["paper_md_list"]
        for segement_res in self._workflow.stream({
            "paper_md_list": paper_md_list,
            "index": 0,
            "paper_info": paper_info,
            "abstract_summary": abstract_summary,
            "next_tip": By3StepsSplitter.first_tip
            }, {"recursion_limit": 999}, stream_mode="custom"):
            yield segement_res["text_segment"]
        
    def _text_generator(self, state: SplitState):
        text_index = state["index"]
        paper_md_list = state["paper_md_list"]
        if text_index >= len(paper_md_list):
            cout(f'\n_text_generator:\nindex: {text_index+1}\nnext_or_end: END\n')
            return {"next_or_end": "END"}
        cout(f'\n_text_generator:\nindex: {text_index+1}\ntext: {paper_md_list[text_index][:50]}\nnext_or_end: NEXT\n')
        return {"index": text_index+1, "text": paper_md_list[text_index], "next_or_end": "NEXT"}
    
    def _next_route(self, state: SplitState):
        cout(f'\n_next_route:\nindex: {state["index"]}\nnext_or_end: {state["next_or_end"]}\n')
        return state["next_or_end"]
        
    def _paper_split_adviser(self, state: SplitState):
        """第一步 —— 文本段结构分析与分割建议生成"""
        user_content = pre_adviser_prompt.invoke({
            "paper_info": state["paper_info"],
            "research_topic": state["abstract_summary"].get("research_topic", "没有相关信息"),
            "next_tip": state["next_tip"],
            "text": state["text"]
            }).messages[0].content
        
        resp = self._llm.invoke([{
                "role": "system",
                "content": pre_adviser_system_msg,
            },{
                "role": "user",
                "content": user_content,
            }]
        )
        split_tip = resp.content
        cout(f'\n_paper_split_adviser:\nindex: {state["index"]}\nsplit_tip: {split_tip}\n')
        return {"split_tip": split_tip}
    
    def _paper_text_splitter(self, state: SplitState):
        """第二步 —— 文本分割与公式 / 表格纠正"""
        user_content = splitter_prompt.invoke({
            "paper_info": state["paper_info"],
            "research_topic": state["abstract_summary"].get("research_topic", "没有相关信息"),
            "split_tip": state["split_tip"],
            "text": state["text"]}).messages[0].content
        
        resp = self._llm.invoke([{
                "role": "system",
                "content": splitter_system_msg,
            },{
                "role": "user",
                "content": user_content,
            }]
        )
        split_res = resp.content
        
        writer = get_stream_writer()
        segments = (s.strip() for s in split_res.split("---shenll---"))
        i = 0
        for s in segments:
            if len(s) == 0:
                continue
            if len(s) >= self._max_chunk_size:
                for s_chunk in self._text_splitter.split_text(s):
                    writer({"text_segment": s_chunk})
            else:
                writer({"text_segment": s})
            i = i + 1
        cout(f'\n_paper_text_splitter:\nindex: {state["index"]}\nsegments len: {i}\ntotal split_res len: {len(split_res)}\n')
        return {"split_res": split_res}
    
    def _next_read_adviser(self, state: SplitState):
        """第三步 —— 下一段分割铺垫信息提取"""
        user_content = next_adviser_prompt.invoke({
            "paper_info": state["paper_info"], 
            "research_topic": state["abstract_summary"].get("research_topic", "没有相关信息"),
            "core_method": state["abstract_summary"].get("core_method", "没有相关信息"),
            "research_goal": state["abstract_summary"].get("research_goal", "没有相关信息"),
            "split_tip": state["split_tip"],
            "split_res": state["split_res"]}).messages[0].content
        
        resp = self._llm.invoke([{
                "role": "system",
                "content": next_adviser_system_msg,
            },{
                "role": "user",
                "content": user_content,
            }]
        )
        next_tip = resp.content
        cout(f'\n_next_read_adviser:\nindex: {state["index"]}\nnext_tip: {next_tip}\n')
        return {"next_tip": next_tip}
    
    def _build_graph(self):
        self._graph = StateGraph(SplitState)

        # Add the nodes
        self._graph.add_node("text_generator", self._text_generator)
        self._graph.add_node("paper_split_adviser", self._paper_split_adviser)
        self._graph.add_node("paper_text_splitter", self._paper_text_splitter)
        self._graph.add_node("next_read_adviser", self._next_read_adviser)

        # Add edges to connect nodes
        self._graph.add_edge(START, "text_generator")
        self._graph.add_conditional_edges(
            "text_generator",
            self._next_route,
            {
                "END": END,
                "NEXT": "paper_split_adviser",
            },
        )
        self._graph.add_edge("paper_split_adviser", "paper_text_splitter")
        self._graph.add_edge("paper_text_splitter", "next_read_adviser")
        self._graph.add_edge("next_read_adviser", "text_generator")
        
        # Compile the workflow
        self._workflow = self._graph.compile()